import torch
from torch import nn
from torch.utils.tensorboard import SummaryWriter
from PIL import Image
from torchvision import transforms

# md_img/study_conv.png
input = torch.tensor([
    [1, 2, 0, 3, 1],
    [0, 1, 2, 3, 1],
    [1, 2, 1, 0, 0],
    [5, 2, 3, 1, 1],
    [2, 1, 0, 1, 1],
])
'''
input = torch.tensor([
    [1, 2, 0, 3, 1],
    [0, 1, 2, 3, 1],
    [1, 2, 1, 0, 0],
    [5, 2, 3, 1, 1],
    [2, 1, 0, 1, 1],
],dtype=torch.float)
'''

input = torch.reshape(input, (-1, 1, 5, 5))
print(input)


class MyMod(nn.Module):
    def __init__(self):
        super(MyMod, self).__init__()
        self.maxpool1 = nn.MaxPool2d(kernel_size=3, ceil_mode=True)

    def forward(self, input):
        output = self.maxpool1(input)
        return output


mymod = MyMod()
output = mymod(input)
print(output)

# 一张图片经过最大池化效果
# 图片变模糊了，尺寸变小
writer = SummaryWriter("../logs")
image_path = "data/train/ants_image/0013035.jpg"
img_PIL = Image.open(image_path)
# TOTensor
trans_totensor = transforms.ToTensor()
tensor_img = trans_totensor(img_PIL)
print(tensor_img.shape)
writer.add_image("ToTensor", tensor_img, global_step=1)

mm = MyMod()
output = mm(tensor_img)
print(output.shape)
writer.add_image("ToTensor", output, global_step=2)


writer.close()